/******************************************************************************
	Cobalt - utility libraries for building XNA games
	Copyright (C) 2008 Philip Djonov

	This program is free software; you can redistribute it and/or modify it
	under the terms of the GNU General Public License as published by the Free
	Software Foundation; either version 2 of the License, or (at your option)
	any later version.

	This program is distributed in the hope that it will be useful, but
	WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
	or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
	for more details.

	You should have received a copy of the GNU General Public License along
	with this program; if not, write to the
	
		Free Software Foundation, Inc.
		51 Franklin Street, Fifth Floor
		Boston, MA  02110-1301, USA.
******************************************************************************/

using System;
using System.Collections.Generic;

using Microsoft.Xna.Framework;
using Microsoft.Xna.Framework.Content;
using Microsoft.Xna.Framework.Graphics;

namespace Cobalt.Framework.Graphics
{
	public class LayerShader
	{
		internal int sortOrder;

		private MaterialManager owner;
		internal Effect effect;
		private EffectTechnique technique;

		internal int sortOverride;
		internal LayerShader cloneSource;

		private int resetPass;

		private int shadingRank;
		private List<EffectTechnique> validTechniques;

		internal LayerShader( MaterialManager owner, Effect effect )
		{
			this.owner = owner;
			this.effect = effect;

			validTechniques = new List<EffectTechnique>( effect.Techniques.GetValidTechniques() );
			SetTechnique( validTechniques[0] );

			foreach( var tech in effect.Techniques )
			{
				EffectAnnotation annot = technique.Annotations["Sort"];
				if( annot != null )
				{
					sortOverride = annot.GetValueInt32();
					break;
				}
			}
		}

		/// <summary>
		/// Gets a value indicating whether this material supports drawing
		/// extra geometry without adverse effects.
		/// </summary>
		/// <remarks>
		/// Certain geometry sources may be able to improve performace by
		/// drawing small amounts of extra geometry not otherwise in the
		/// draw list - if the current material supports that, that is.
		/// </remarks>
		public bool AllowExtraDraws
		{
			get { return true; }
		}

		private void SetTechnique( EffectTechnique technique )
		{
			resetPass = -1;
			for( int i = 0; i < technique.Passes.Count; i++ )
			{
				EffectAnnotation annot = technique.Passes[i].Annotations["IsResetPass"];
				if( annot != null && annot.GetValueBoolean() )
				{
					resetPass = i;
					break;
				}
			}

			effect.CurrentTechnique = technique;
			shadingRank = validTechniques.IndexOf( technique );

			this.technique = technique;
		}

		public MaterialManager Owner
		{
			get { return owner; }
		}

		public virtual int Begin( ShadingContext context )
		{
			if( bindingsDirty )
				UpdateBindings();

			int targetRank = Helpers.Clamp( context.ShadingComplexityRank, 0, validTechniques.Count - 1 );
			if( targetRank != shadingRank )
				SetTechnique( validTechniques[targetRank] );

			owner.UpdateSharedParameters( context );

			effect.Begin();

			if( resetPass != -1 )
				return resetPass;

			return technique.Passes.Count;
		}
		public virtual void End( ShadingContext context )
		{
			if( resetPass != -1 )
			{
				technique.Passes[resetPass].Begin();
				technique.Passes[resetPass].End();
			}

			effect.End();
		}

		public virtual void BeginPass( int pass, ShadingContext context )
		{
			technique.Passes[pass].Begin();
		}

		public virtual void EndPass( int pass, ShadingContext context )
		{
			technique.Passes[pass].End();
		}

		private bool bindingsDirty = true;

		internal void UpdateBindings()
		{
			foreach( EffectParameter param in effect.Parameters )
			{
				EffectAnnotation annot = param.Annotations["BindTexDimension"];
				if( annot == null || annot.ParameterType != EffectParameterType.String )
					continue;

				EffectParameter texParam = effect.Parameters[annot.GetValueString()];
				if( texParam == null )
					continue;

				Vector4 value = Vector4.Zero;
				switch( texParam.ParameterType )
				{
				case EffectParameterType.Texture:
					if( param.ColumnCount == 1 )
						goto case EffectParameterType.Texture1D;
					if( param.ColumnCount == 2 )
						goto case EffectParameterType.Texture2D;
					if( param.ColumnCount == 3 )
						goto case EffectParameterType.Texture3D;
					break;

				case EffectParameterType.Texture1D:
				case EffectParameterType.Texture2D:
					{
						Texture2D tex = texParam.GetValueTexture2D();
						if( tex != null )
							value = new Vector4( tex.Width, tex.Height, 0, 1 );
					}
					break;

				case EffectParameterType.Texture3D:
					{
						Texture3D tex = texParam.GetValueTexture3D();
						if( tex != null )
							value = new Vector4( tex.Width, tex.Height, tex.Depth, 1 );
					}
					break;

				case EffectParameterType.TextureCube:
					{
						TextureCube tex = texParam.GetValueTextureCube();
						if( tex != null )
							value = new Vector4( tex.Size, tex.Size, 0, 1 );
					}
					break;
				}

				switch( param.ColumnCount )
				{
				case 1:
					param.SetValue( value.X );
					break;

				case 2:
					param.SetValue( new Vector2( value.X, value.Y ) );
					break;

				case 3:
					param.SetValue( new Vector3( value.X, value.Y, value.Z ) );
					break;

				case 4:
					param.SetValue( value );
					break;
				}
			}

			bindingsDirty = false;
		}

		internal void SetParameter( string parameterName, Texture value )
		{
			EffectParameter param = effect.Parameters[parameterName];
			if( param == null )
				return;

			param.SetValue( value );
			bindingsDirty = true;
		}

		internal LayerShader Clone()
		{
			Effect clonedFx = effect.Clone( effect.GraphicsDevice );

			LayerShader ret = owner.GetLayer( clonedFx );
			ret.bindingsDirty = false;

			return ret;
		}
	}

	public sealed class Material : IDisposable
	{
		private MaterialManager owner;
		internal Material( MaterialManager owner )
		{
			this.owner = owner;
		}

		internal void Loaded()
		{
			owner.RegisterMaterial( this );
		}

		public MaterialManager Owner { get { return owner; } }

		private string name;
		public string Name
		{
			get { return name; }
			internal set { name = value; }
		}

		void IDisposable.Dispose()
		{
			owner.UnregisterMaterial( this );
		}

		public Material Clone()
		{
			Material ret = new Material( owner );

			ret.tags = tags.Clone();

			for( int i = 0; i < layers.Length; i++ )
			{
				if( layers[i] != null )
					ret.layers[i] = layers[i].Clone();
			}

			ret.Loaded();

			return ret;
		}

		internal int sortOrder;

		public void SetParameter( string parameterName, Texture value )
		{
			foreach( LayerShader sh in layers )
			{
				if( sh == null )
					continue;

				sh.SetParameter( parameterName, value );
			}
		}

		internal TagsDictionary tags;
		public TagsDictionary Tags { get { return tags; } }

		public enum Layer
		{
			Depth,
			Light,
		}

		internal const int NumLayers = (int)Layer.Light + 1;

		internal static bool IsValidLayer( Layer layer )
		{
			return (int)layer >= 0 && (int)layer < NumLayers;
		}

		internal LayerShader[] layers = new LayerShader[NumLayers];

		public LayerShader GetLayer( Layer layer )
		{
			switch( layer )
			{
				case Layer.Depth:
				case Layer.Light:
					return layers[(int)layer];

				default:
					throw new ArgumentException( "Unknown layer type." );
			}
		}

		public LayerShader DepthShader
		{
			get { return layers[(int)Layer.Depth]; }
			private set { layers[(int)Layer.Depth] = value; }
		}

		public LayerShader LightShader
		{
			get { return layers[(int)Layer.Light]; }
			private set { layers[(int)Layer.Light] = value; }
		}
	}

	public interface IMaterialService
	{
		MaterialManager MaterialManager { get; }
	}

	public class MaterialManager : IMaterialService, IDisposable
	{
		public MaterialManager( IServiceProvider services, ContentManager contentManager )
		{
			if( services == null )
				throw new ArgumentNullException();
			if( contentManager == null )
				throw new ArgumentNullException();

			graphics = services.GetService( typeof( IGraphicsDeviceService ) )
				as IGraphicsDeviceService;

			if( graphics == null )
				throw new ArgumentException( "Missing IGraphicsDeviceService on ContentManager." );

			this.contentManager = contentManager;
		}

		private string rootDirectory = string.Empty;
		public string RootDirectory
		{
			get { return rootDirectory; }
			set { rootDirectory = value; }
		}

		private IGraphicsDeviceService graphics;

		#region Registration and Layer Sorting
		private List<Material> materials = new List<Material>();
		internal int[][] layerSortIndices = new int[Material.NumLayers][];

		private void SortLayer( Material.Layer layer )
		{
			List<LayerShader> layers = new List<LayerShader>( materials.Count );
			foreach( Material mat in materials )
			{
				LayerShader sh = mat.layers[(int)layer];
				if( sh == null || layers.Contains( sh ) )
					continue;

				layers.Add( sh );
			}

			layers.Sort( delegate( LayerShader a, LayerShader b )
				{
					if( a == null )
					{
						if( b == null )
							return 0;

						return 1;
					}
					else
					{
						if( b == null )
							return -1;
					}

					return a.sortOverride - b.sortOverride;
				} );

			for( int i = 0; i < layers.Count; i++ )
				layers[i].sortOrder = i;

			Helpers.EnsureArray( ref layerSortIndices[(int)layer], materials.Count );

			for( int i = 0; i < materials.Count; i++ )
			{
				LayerShader sh = materials[i].layers[(int)layer];
				layerSortIndices[(int)layer][i] = sh != null ? sh.sortOrder : -1;
			}
		}

		private int updateCount = 0;
		private void SortMaterials()
		{
			if( updateCount > 0 )
				return;

#warning This can mess up any in-progress DrawLists, find a way to either fix them up in flight or throw an exception.

			for( int i = 0; i < materials.Count; i++ )
				materials[i].sortOrder = i;

			for( int i = 0; i < Material.NumLayers; i++ )
				SortLayer( (Material.Layer)i );
		}

		public void BeginUpdate()
		{
			updateCount++;
		}

		public void EndUpdate()
		{
			if( updateCount == 0 )
				throw new InvalidOperationException();

			updateCount--;

			if( updateCount != 0 )
				return;

			SortMaterials();
			RequireSharedEffect( materials.Count > 0 ? materials[0] : null );
		}

		internal void RegisterMaterial( Material material )
		{
			if( updateCount == 0 )
				RequireSharedEffect( material );

			materials.Add( material );
			SortMaterials();
		}

		internal void UnregisterMaterial( Material material )
		{
			foreach( LayerShader layer in material.layers )
			{
				if( layer != null && layer.effect == sharedEffect )
				{
					UnloadSharedParams();
					sharedEffect = null;
				}
			}

			materials.Remove( material );
			SortMaterials();
		}
		#endregion

		#region Shared Parameter Handling
		private Effect sharedEffect;
		private EffectParameter ehT, ehM, ehMv, ehMvp, ehWCam, ehMCam, ehDepthRange;

		private bool LoadSharedParams( Effect effect )
		{
			ehT = effect.Parameters.GetParameterBySemantic( "TIME" );
			ehM = effect.Parameters.GetParameterBySemantic( "MODEL" );
			ehMv = effect.Parameters.GetParameterBySemantic( "MODELVIEW" );
			ehMvp = effect.Parameters.GetParameterBySemantic( "MODELVIEWPROJECTION" );
			ehWCam = effect.Parameters.GetParameterBySemantic( "WORLDCAMERAPOSITION" );
			ehMCam = effect.Parameters.GetParameterBySemantic( "MODELCAMERAPOSITION" );
			ehDepthRange = effect.Parameters.GetParameterBySemantic( "DEPTHRANGE" );

			return ehM != null && ehMv != null && ehMvp != null && ehMCam != null && ehWCam != null && ehDepthRange != null;
		}

		private void UnloadSharedParams()
		{
			ehT = null;
			ehM = null;
			ehMv = null;
			ehMvp = null;
			ehWCam = null;
			ehMCam = null;
			ehDepthRange = null;
		}

		private bool loadingShared = false;
		private void RequireSharedEffect( Material existingInstance )
		{
			if( sharedEffect != null || loadingShared )
				return;

			if( existingInstance != null )
			{
				foreach( LayerShader layer in existingInstance.layers )
					if( layer != null && LoadSharedParams( layer.effect ) )
					{
						sharedEffect = layer.effect;
						return;
					}
			}

			loadingShared = true;
			Effect fx = null;

			try
			{
				fx = contentManager.Load<Effect>( GetContentPath( @"FX\Shared" ) );
			}
			catch( ContentLoadException ex )
			{
				throw new InvalidOperationException( "Unable to load shared effect.", ex );
			}

			if( LoadSharedParams( fx ) )
				sharedEffect = fx;
			loadingShared = false;

			if( sharedEffect == null )
				throw new InvalidOperationException( "Unable to load shared effect." );
		}

		internal void UpdateSharedParameters( ShadingContext context )
		{
			ehT.SetValue( context.RenderTime );
			ehM.SetValue( context.Model );
			ehMv.SetValue( context.ModelView );
			ehMvp.SetValue( context.ModelViewProjection );
			ehWCam.SetValue( context.WorldCameraPosition );
			ehMCam.SetValue( context.ModelCameraPosition );
			ehDepthRange.SetValue( context.DepthRange );
		}
		#endregion

		#region Loading
		private ContentManager contentManager;

		private Dictionary<Effect, LayerShader> layerShaders = new Dictionary<Effect, LayerShader>();
		internal LayerShader GetLayer( Effect effect )
		{
			LayerShader ret;
			if( !layerShaders.TryGetValue( effect, out ret ) )
			{
				ret = new LayerShader( this, effect );
				layerShaders.Add( effect, ret );
			}
			return ret;
		}

		private string GetContentPath( string name )
		{
			if( rootDirectory != null )
				name = System.IO.Path.Combine( rootDirectory, name );

			return name;
		}


		public Material GetMaterial( string name )
		{
			string matName = GetContentPath( name );

			try
			{
				return contentManager.Load<Material>( matName );
			}
			catch( ContentLoadException )
			{
				System.Diagnostics.Debug.WriteLine( string.Format( "Can't find material '{0}' at '{1}' - loading 'Default' instead.", name, matName ) );
				return contentManager.Load<Material>( GetContentPath( "Default" ) );
			}
		}

		public ContentManager ContentManager
		{
			get { return contentManager; }
		}
		#endregion

		MaterialManager IMaterialService.MaterialManager
		{
			get { return this; }
		}

		#region IDisposable Members

		public void Dispose()
		{
			if( graphics == null )
				return;

			materials.Clear();

			graphics = null;
		}

		#endregion
	}
}
